#include <zephyr.h>

#include <stdio.h>
#include <string.h>

#include "mbedtls/rsa.h"
#include "mbedtls/cipher.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/platform.h"

#define assert_exit(cond, ret) \
    do { if (!(cond)) { \
        printf("  !. assert: failed [line: %d, error: -0x%04X]\n", __LINE__, -ret); \
        goto cleanup; \
    } } while (0)

typedef struct cipher
{
  uint8_t cipher_pre[2048/8];
  uint8_t cipher_back[2048/8];
}cipher_str;  
 
 
typedef struct msg_sign
{
  uint8_t sign_pre[2048/8];
  uint8_t mag_back[2048/8];
}msg_sign_str;   


cipher_str cipher,cipher_bak;
msg_hash_str mas_hash,mas_hash_bak;


  

 uint8_t key[32] =
{
    0x06, 0xa9, 0x21, 0x40, 0x36, 0xb8, 0xa1, 0x5b,
    0x51, 0x2e, 0x03, 0xd5, 0x34, 0x12, 0x00, 0x06
};

uint8_t iv[32] =
{
    0x3d, 0xaf, 0xba, 0x42, 0x9d, 0x9e, 0xb4, 0x30,
    0xb4, 0x22, 0xda, 0x80, 0x2c, 0x9f, 0xac, 0x41
};



static void dump_buf(char *info, uint8_t *buf, uint32_t len)
{
    mbedtls_printf("%s", info);
    for (int i = 0; i < len; i++) {
        mbedtls_printf("%s%02X%s", i % 16 == 0 ? "\n     ":" ", 
                        buf[i], i == len - 1 ? "\n":"");
    }
}

static int entropy_source(void *data, uint8_t *output, size_t len, size_t *olen)
{
    uint32_t seed;

    seed = sys_rand32_get();
    if (len > sizeof(seed)) {
        len = sizeof(seed);
    }

    memcpy(output, &seed, len);

    *olen = len;
    return 0;
}

void cipher_decrypt(int type,uint8_t *input,uint8_t *output)
{
    size_t len;
    int olen = 0;

    mbedtls_cipher_context_t ctx;
    const mbedtls_cipher_info_t *info;

    mbedtls_platform_set_printf(printf);

    mbedtls_cipher_init(&ctx);
    info = mbedtls_cipher_info_from_type(type);

    mbedtls_cipher_setup(&ctx, info);
                   
    //解密
    mbedtls_cipher_setkey(&ctx,key, sizeof(key)*8, MBEDTLS_DECRYPT);
    
    mbedtls_cipher_crypt(&ctx,iv,sizeof(iv),input, strlen(input),output, &len );     
    olen += len;

    mbedtls_cipher_finish(&ctx, output + len, &len);
    olen += len;
    
        
    printf("\n %s",output);
    dump_buf("\n  cipher aes decrypt:", output, olen);

    mbedtls_cipher_free(&ctx);
}


void cipher_encrypt(int type,uint8_t *input,uint8_t *output)
{
    size_t len;
    int olen = 0;

    mbedtls_cipher_context_t ctx;
    const mbedtls_cipher_info_t *info;

    mbedtls_platform_set_printf(printf);

    mbedtls_cipher_init(&ctx);
    info = mbedtls_cipher_info_from_type(type);

    mbedtls_cipher_setup(&ctx, info);


    mbedtls_cipher_setkey(&ctx, key, sizeof(key)*8, MBEDTLS_ENCRYPT);
    mbedtls_cipher_update(&ctx, input, strlen(input), output, &len);
    olen += len;

    mbedtls_cipher_finish(&ctx, output + len, &len);
    olen += len;
    

   mbedtls_printf("\n 3. AES密文写入成功\n");
    
    

    dump_buf("\n  cipher aes encrypt:",output, olen);

    mbedtls_cipher_free(&ctx);
}



int main(void)
{
    int ret;
    size_t olen = 0;
    uint8_t out[2048/8];
    uint8_t sig[2048/8];

    uint16_t s_len = 0;

    mbedtls_platform_set_printf(printf);
    mbedtls_platform_set_snprintf(snprintf);

    mbedtls_rsa_context ctx;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    const char *pers = "simple_rsa";
    const char *msg = "Hello, World!";

    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_rsa_init(&ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);
    
    mbedtls_entropy_add_source(&entropy, entropy_source, NULL,
                               MBEDTLS_ENTROPY_MAX_GATHER,
                               MBEDTLS_ENTROPY_SOURCE_STRONG);
    ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy, 
                                    (const uint8_t *) pers, strlen(pers));
    assert_exit(ret == 0, ret);
    mbedtls_printf("\n  . setup rng ... ok\n");
    
    
    

    mbedtls_printf("\n  ! RSA Generating large primes may take minutes! \n");
    ret = mbedtls_rsa_gen_key(&ctx, mbedtls_ctr_drbg_random, 
                                        &ctr_drbg, 2048, 65537);
    assert_exit(ret == 0, ret);                                    
    mbedtls_printf("\n  1. RSA generate key ... ok\n");

    
    
    
    
    //1.将消息哈希并用公钥签名     H=Hash(M)   S=Sig(H,Kda)
    ret = mbedtls_rsa_pkcs1_sign(&ctx, mbedtls_ctr_drbg_random, &ctr_drbg, 
                                    MBEDTLS_RSA_PRIVATE, MBEDTLS_MD_SHA256,
                                    sizeof(msg), msg, sig);
    assert_exit(ret == 0, ret);
    dump_buf("  2. 将消息哈希并签名:", sig, sizeof(sig));
    
    //2.将签名和消息写入文件     M//S
    FILE *fptr;
   fptr = fopen("msg_hash.txt", "wb");
   if(fptr == NULL)
   {
      printf("Error!");
      exit(1);
   } 
   fprintf(fptr,"%s", msg);
   fprintf(fptr,"%s", "0");//分割符号
   fprintf(fptr,"%s", "0");
   fprintf(fptr,"%s", "0");
   fprintf(fptr,"%s", sig);
   fclose(fptr);
   
   mbedtls_printf("\n 签名和消息写入成功\n");
   
   
    //3.读取msg_hash.txt并AES加密   C = E（M//S，K）
    FILE *fp;
    fp = fopen("msg_hash.txt" , "r");
    fseek( fp , 0 , SEEK_END );
    int file_size;
    file_size = ftell( fp );
    printf( "%d" , file_size );
    char *msg_hash;
    fseek( fp , 0 , SEEK_SET);
    msg_hash =  (char *)malloc( file_size * sizeof( char ) );
    fread( msg_hash , file_size , sizeof(char) , fp);
    mbedtls_printf("\n 读取msg_hash成功\n");
    
   //将密文写入文件 
   FILE *f_cipher;
   f_cipher = fopen("cipher.txt", "wb");
   if(f_cipher == NULL)
   {
      printf("Error!");
      exit(1);
   } 

    
    //AES加密
    cipher_encrypt(MBEDTLS_CIPHER_AES_256_CTR,msg_hash,f_cipher);
    
    

    //4 .RSA加密AES公钥         C1 = E（K，Keb）
    ret = mbedtls_rsa_pkcs1_encrypt(&ctx, mbedtls_ctr_drbg_random, 
                            &ctr_drbg, MBEDTLS_RSA_PUBLIC, strlen(key), key, out);
    assert_exit(ret == 0, ret);                              
    dump_buf("\n  4. RSA 加密AES公钥完成", out, sizeof(out));
    


   //   C // C1
   fprintf(f_cipher,"%s", "0");//分割符号
   fprintf(f_cipher,"%s", "0");
   fprintf(f_cipher,"%s", "0");
   fprintf(f_cipher,"%s", out);

   fclose(f_cipher);
   mbedtls_printf("\n 加密公钥写入成功\n");
   
   
   
  /******************************************************************************************************************/
  /******************************************解码逆操作*************************************************************/
  
    FILE *fr;
    fr = fopen("cipher.txt" , "r");
    fseek( fr , 0 , SEEK_END );
    int f_size;
    f_size = ftell( fr );
    printf( "%d" , file_size );
    char *cipher;
    fseek( fr , 0 , SEEK_SET);
    cipher =  (char *)malloc( f_size * sizeof( char ) );
    fread( cipher , f_size , sizeof(char) , fr);
    dump_buf("\n  提取密文成功:", cipher, sizeof(cipher));
    



    
      //将提取写入文件 
   FILE *fptr1;
   fptr1 = fopen("comp.txt", "wb");
   if(fptr1 == NULL)
   {
      printf("Error!");
      exit(1);
   } 
   fprintf(fptr1,"%s",msg_hash_cipher);
   
   fclose(fptr1);
   mbedtls_printf("\n 写入文件对比\n");
    
    
    //AES解密出信息和哈希
    //cipher_decrypt(MBEDTLS_CIPHER_AES_256_CTR,msg_hash_cipher);
  
  /*
  
   //解密出AES公钥
    ret = mbedtls_rsa_pkcs1_decrypt(&ctx, mbedtls_ctr_drbg_random, &ctr_drbg, 
                                MBEDTLS_RSA_PRIVATE, &olen, out, out, sizeof(out));
    assert_exit(ret == 0, ret);                              
    
    out[olen] = 0;
    mbedtls_printf("\n  3. RSA decryption ... ok\n     %s\n", out);
    
    
    //
    
    FILE *fr2;
    fr2 = fopen("cipher.txt" , "r");
    fseek( fr2 , 0 , SEEK_END );
    int f2_size;
    f2_size = ftell( fr2 );
    printf( "%d" , f2_size );
    char *cipher;
    fseek( fr2 , 0 , SEEK_SET);
    cipher =  (char *)malloc( f2_size * sizeof( char ) );
    fread( cipher, f2_size , sizeof(char) , fr2);
    dump_buf("\n  提取密文成功:", cipher, sizeof(cipher));
    
        
 
    //AES解密出信息和哈希
    cipher_decrypt(MBEDTLS_CIPHER_AES_256_CTR,cipher,key)；
    
    
    FILE *fr3;
    fr3 = fopen("plain.txt" , "r");
    fseek( fr2 , 0 , SEEK_END );
    int f3_size;
    f3_size = ftell( fr3 );
    printf( "%d" , f3_size );
    char *plain;
    fseek( fr3 , 0 , SEEK_SET);
    plain =  (char *)malloc( f3_size * sizeof( char ) );
    fread( plain, f3_size , sizeof(char) , fr3);
    dump_buf("\n  提取明文成功:", plain, sizeof(plain));
    
    
    */
    
    
    //
    
    
    
    
    
    
    


  

    ret = mbedtls_rsa_pkcs1_encrypt(&ctx, mbedtls_ctr_drbg_random, 
                            &ctr_drbg, MBEDTLS_RSA_PUBLIC, strlen(msg), msg, out);
    assert_exit(ret == 0, ret);                              
    dump_buf("\n  保留", out, sizeof(out));


    ret = mbedtls_rsa_pkcs1_decrypt(&ctx, mbedtls_ctr_drbg_random, &ctr_drbg, 
                                MBEDTLS_RSA_PRIVATE, &olen, out, out, sizeof(out));
    assert_exit(ret == 0, ret);                              
    
    out[olen] = 0;
    mbedtls_printf("\n  3. RSA decryption ... ok\n     %s\n", out);
    
    
    
    
    
    

    ret = memcmp(out, msg, olen);
    assert_exit(ret == 0, ret);      
    mbedtls_printf("\n  4. RSA Compare results and plaintext ... ok\n");
    
    
    
    
    
    
    
    
    
    

cleanup:
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);
    mbedtls_rsa_free(&ctx);

    return ret;
}
